import math
import numpy as np
from quaternion import quaternion


def rotation_matrix_z(theta):
    """
    根据绕Z轴的角度theta（以弧度为单位）计算旋转矩阵。
    """
    c, s = np.cos(theta), np.sin(theta)
    return np.array([[c, -s, 0],
                     [s, c, 0],
                     [0, 0, 1]])


def gen_corners(box):
    # boxes: [x,y,z,l,w,h,r]
    w, l, h = box[[4, 3, 5]]
    # 3D bounding box corners. (Convention: x points forward, y to the left, z up.)
    x_corners = l / 2 * np.array([1, 1, 1, 1, -1, -1, -1, -1])
    y_corners = w / 2 * np.array([1, -1, -1, 1, 1, -1, -1, 1])
    z_corners = h / 2 * np.array([1, 1, -1, -1, 1, 1, -1, -1])
    corners = np.vstack((x_corners, y_corners, z_corners))
    # Rotate
    rotation_matrix = rotation_matrix_z(box[6])
    corners = np.dot(rotation_matrix, corners)
    # Translate
    x, y, z = box[[0, 1, 2]]
    corners[0, :] = corners[0, :] + x
    corners[1, :] = corners[1, :] + y
    corners[2, :] = corners[2, :] + z
    return corners


def isRotationMatrix(R):
    Rt = np.transpose(R)
    shouldBeIdentity = np.dot(Rt, R)
    I = np.identity(3, dtype=R.dtype)
    n = np.linalg.norm(I - shouldBeIdentity)
    return n < 1e-6


# rotationMatrixToEulerAngles 用于旋转矩阵转欧拉角
def rotationMatrixToEulerAngles(R):
    assert (isRotationMatrix(R))
    sy = math.sqrt(R[0, 0] * R[0, 0] + R[1, 0] * R[1, 0])
    singular = sy < 1e-6
    if not singular:
        x = math.atan2(R[2, 1], R[2, 2])
        y = math.atan2(-R[2, 0], sy)
        z = math.atan2(R[1, 0], R[0, 0])
    else:
        x = math.atan2(-R[1, 2], R[1, 1])
        y = math.atan2(-R[2, 0], sy)
        z = 0
    return np.array([x, y, z])




def quaternion_to_euler_angles(wxyz:list, axes='sxyz'):
    """
    将四元数转换为欧拉角。

    参数:
    q (quaternion): 一个表示旋转的四元数对象。
    axes (str): 指定旋转顺序的轴，默认为'sxyz'（对应于ZYX旋转顺序，即偏航-俯仰-滚动）。
                可以是'sxyz', 'sxzy', 'syxz', 'syzy', 'szxy', 'szyz'之一。

    返回:
    tuple: 包含三个欧拉角的元组（以弧度为单位）。
    """
    w, x, y, z = wxyz

    # 计算欧拉角
    # 注意：这里的公式可能因旋转顺序和定义的不同而有所变化
    # 下面的公式假设了ZYX旋转顺序（偏航-俯仰-滚动）
    sinr_cosp = 2 * (w * x + y * z)
    cosr_cosp = 1 - 2 * (x ** 2 + y ** 2)
    roll = np.arctan2(sinr_cosp, cosr_cosp)  # 滚动角（绕X轴）

    sinp = 2 * (w * y - z * x)
    if abs(sinp) >= 1:
        pitch = np.copysign(np.pi / 2, sinp)  # 使用90度来避免除以零
    else:
        pitch = np.arcsin(sinp)  # 俯仰角（绕Y轴）

    siny_cosp = 2 * (w * z + x * y)
    cosy_cosp = 1 - 2 * (y ** 2 + z ** 2)
    yaw = np.arctan2(siny_cosp, cosy_cosp)  # 偏航角（绕Z轴）

    # 根据指定的轴顺序调整角度
    if axes == 'sxyz':
        pass  # ZYX顺序（默认）
    elif axes == 'sxzy':
        yaw, pitch, roll = yaw, roll, -pitch
    elif axes == 'syxz':
        yaw, pitch, roll = -yaw, pitch, roll
    elif axes == 'syzy':
        yaw, pitch, roll = -yaw, -pitch, roll
    elif axes == 'szxy':
        yaw, pitch, roll = yaw, -pitch, -roll
    elif axes == 'szyz':
        yaw, pitch, roll = -yaw, -pitch, -roll
    else:
        raise ValueError("axes must be one of 'sxyz', 'sxzy', 'syxz', 'syzy', 'szxy', 'szyz'")

    # 返回以弧度为单位的欧拉角
    return roll, pitch, yaw


# # 示例用法
# roll, pitch, yaw = quaternion_to_euler_angles([0.707, 0, 0, 0.707], axes='sxyz')
# print(f"Roll: {roll}, Pitch: {pitch}, Yaw: {yaw}")

def quaternion_to_rotation_matrix(w, x, y, z):
    R = np.array([
        [1 - 2 * y ** 2 - 2 * z ** 2, 2 * x * y - 2 * z * w, 2 * x * z + 2 * y * w],
        [2 * x * y + 2 * z * w, 1 - 2 * x ** 2 - 2 * z ** 2, 2 * y * z - 2 * x * w],
        [2 * x * z - 2 * y * w, 2 * y * z + 2 * x * w, 1 - 2 * x ** 2 - 2 * y ** 2]
    ])
    return R


def pose2mat(pose):
    rt = np.eye(4)
    rt[:3, :3] = quaternion_to_rotation_matrix(pose.orientation.w, pose.orientation.x, pose.orientation.y,
                                               pose.orientation.z)
    rt[:3, 3] = np.array([pose.position.x, pose.position.y, pose.position.z])
    return rt


def relative_pose(ego, obj):
    ego2local = pose2mat(ego)
    obj2local = pose2mat(obj)
    obj2eog = np.linalg.inv(ego2local) @ obj2local

    _, _, yaw = rotationMatrixToEulerAngles(obj2eog[:3, :3])
    # 组合成相对位姿格式
    relative_pose_info = {
        'dx': obj2eog[0, 3],
        'dy': obj2eog[1, 3],
        'dz': obj2eog[2, 3],
        'dyaw': yaw
    }

    return relative_pose_info
